package com.tool.sqlexecutor.util;

import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.tool.sqlexecutor.config.DataSourceProperty;
import com.tool.sqlexecutor.config.SSH;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.jdbc.datasource.DriverManagerDataSource;

import javax.sql.DataSource;

public class SSHClient {
    private static final Log logger = LogFactory.getLog(SSHClient.class);
    private static final String JDBC_DRIVER = "com.mysql.cj.jdbc.Driver";
    private SSH ssh;
    private DataSourceProperty dataSourceProperty;
    private Session session;
    private int port;

    public SSHClient(SSH ssh, DataSourceProperty dataSourceProperty) {
        this.ssh = ssh;
        this.dataSourceProperty = dataSourceProperty;
    }

    public void start() throws JSchException {
        if (!isEnabled()) {
            return;
        }
        logger.info("ssh is enabled. database:" + dataSourceProperty.getDatabase());

        session = new JSch().getSession(ssh.getUser(), ssh.getHost(), 22);
        session.setPassword(ssh.getPassword());
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect();
        logger.info("ssh session connected! database:" + dataSourceProperty.getDatabase());
        port = session.setPortForwardingL(0, dataSourceProperty.getHost(), dataSourceProperty.getPort());
    }

    public DataSource getDataSource() {
        if (!isEnabled()) {
            return DataSourceBuilder.create()
                    .type(DriverManagerDataSource.class)
                    .driverClassName(JDBC_DRIVER)
                    .url(dataSourceProperty.getUrl())
                    .username(dataSourceProperty.getUser())
                    .password(dataSourceProperty.getPassword())
                    .build();
        }
        return DataSourceBuilder.create()
                .type(DriverManagerDataSource.class)
                .driverClassName(JDBC_DRIVER)
                .url("jdbc:mysql://localhost:" + port + "/" + dataSourceProperty.getDatabase() + "?useUnicode=true&characterEncoding=utf8")
                .username(dataSourceProperty.getUser())
                .password(dataSourceProperty.getPassword())
                .build();
    }

    public void end() {
        if (!isEnabled()) {
            return;
        }
        session.disconnect();
    }

    private boolean isEnabled() {
        return ssh != null && ssh.isEnable();
    }
}
